/*
    Author: Christopher Diggins
    License: MIT Licence 1.0
    Website: http://www.cdiggins.com
*/

#ifndef CVML_PREFIX_TREE_HPP
#define CVML_PREFIX_TREE_HPP

/// Used to find common non-overlapping prefixes. 
struct PrefixTree
{
    struct Node
    {
        std::vector<InstrSeq> seqs;
        std::map<int, Node*> children;
        void* data;

        typedef std::map<int, Node*>::iterator iterator;

        Node() 
            : data(NULL)
        { }

        void GetLeaves(std::vector<Node*>& vec) {
            if (children.size() == 0)
                vec.push_back(this);
            else
            for (iterator i = children.begin(); i != children.end(); i++)
                i->second->GetLeaves(vec);
        }
    };

    Node* root;
    int maxlen;

    PrefixTree(int _maxlen = -1) 
        : root(new Node()), maxlen(_maxlen)
    {
    }

    ~PrefixTree()
    {
        delete root;
    }

    void AddSeqFrom(Instruction* i, void* data = NULL) {
        if (maxlen >= 0) 
            if (!i->LenGtEq(maxlen))
                return;
        InstrSeq seq(i, maxlen);
        
        // If we are using a prefix tree, we should first strip all no-ops.
        Assert(!seq.ContainsOp(op_noop));

        if (maxlen >= 0)
            Assert(seq.Len() == maxlen);

        AddSeq(seq, data);
    }

    void AddSeq(InstrSeq seq, void* data = NULL) {
        AddSeqTo(root, seq, seq.begin, data);
    }

    void AddSeqTo(Node* node, InstrSeq seq, Instruction* cur, void* data = NULL) {
        if (maxlen >= 0) 
            Assert(seq.Len() == maxlen);

        while (cur != seq.end) {
            node = GetOrCreateChild(node, cur);
            cur = cur->next;
        }

        // If we are using a prefix tree, we should have stripped all no-ops.
        Assert(!seq.ContainsOp(op_noop));

        node->seqs.push_back(seq);
        node->data = data;
        return;
    }

    static int GetHashCode(Instruction* i) 
    {
        if (i == NULL)
            return 0;
        return (i->GetHashCode());
    }

    bool NodeHasChild(Node* node, Instruction* i) {
        return MapContains(node->children, GetHashCode(i));
    }

    Node* GetChild(Node* node, Instruction* i) {
        if (!NodeHasChild(node, i))
            return false;
        return node->children[GetHashCode(i)];
    }
        
    Node* GetOrCreateChild(Node* node, Instruction* i) {
        if (!NodeHasChild(node, i)) {
            AddChild(node, i);
        }
        return GetChild(node, i);
    }

    void AddChild(Node* parent, Instruction* cur) {
        Node* child = new Node();
        parent->children[GetHashCode(cur)] = child;
    }

    Node* GetNode(InstrSeq seq) {
        Instruction* i = seq.begin;
        Node* cur = root;
        while (i != seq.end) {
            if (i == NULL) 
                return NULL;
            cur = GetChild(cur, i);
            if (cur == NULL) 
                return NULL;
            i = i->next;
        }
        return cur;
    }

    bool ContainsSeq(InstrSeq seq) {
        return GetNode(seq) != NULL;
    }

    void* GetAssociatedData(InstrSeq seq) {
        Node* node = GetNode(seq);
        if (node == NULL)
            return NULL;
        return node->data;
    }

    void GetLeaves(std::vector<Node*>& vec) {
        root->GetLeaves(vec);
    }
};

#endif
